package com.inspur.edp.ai.flow.node;

import com.fasterxml.jackson.databind.node.ArrayNode;
import com.inspur.edp.ai.flow.engine.FlowContext;
import com.inspur.edp.ai.flow.engine.strategy.DefaultFlowStrategy;
import com.inspur.edp.ai.flow.engine.strategy.LoopBodyFlowStrategy;
import com.inspur.edp.ai.flow.exception.FlowException;
import com.inspur.edp.ai.flow.metadata.AbsNode;
import com.inspur.edp.ai.flow.spi.NodeExecuteResult;
import com.inspur.edp.ai.flow.spi.NodeExecutor;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author lizhaorui
 * @date 2025/8/15
 * @description
 */
public class LoopNodeExecutor implements NodeExecutor {

    @Override
    public String getKind() {
        return LoopNode.KIND;
    }

    @Override
    public NodeExecuteResult execute(FlowContext context, AbsNode node, Map<String, Object> args) {
        LoopNode loopNode = (LoopNode) node;
        Map<String, Object> nodeParameterValue = (Map<String, Object>) args.get(loopNode.getCode());
        int index = 0;
        if (nodeParameterValue != null) {
            Object indexObject = nodeParameterValue.get(LoopNode.ITEM_INDEX_KEY);
            if (indexObject != null) {
                index = (int) indexObject;
            }
        }
        Object iterableExprValue = loopNode.getIterableExprEvaluator().evaluate(args);
        //迭代的值为null，则无需循环
        if (iterableExprValue == null) {
            return this.continueNextNode(loopNode);
        } else {
            Object itemValue = null;
            if (iterableExprValue instanceof List) {
                List list = (List) iterableExprValue;
                if (list.size() > index) {
                    itemValue = list.get(index);
                } else {
                    return this.continueNextNode(loopNode);
                }
            } else if (iterableExprValue instanceof ArrayNode) {
                ArrayNode arrayNode = (ArrayNode) iterableExprValue;
                if (arrayNode.size() > index) {
                    itemValue = arrayNode.get(index);
                } else {
                    return this.continueNextNode(loopNode);
                }
            } else {
                throw new FlowException("not support type " + iterableExprValue.getClass().getName());
            }
            int nextIndex = index + 1;
            return this.executeLoopBody(loopNode, itemValue, nextIndex);
        }
    }

    private NodeExecuteResult executeLoopBody(LoopNode loopNode, Object iterableValue, int index) {
        //开始子流程策略
        LoopBodyFlowStrategy flowStrategy = new LoopBodyFlowStrategy();
        flowStrategy.setNodeId(loopNode.getId());
        flowStrategy.setNodes(loopNode.getNodes());
        flowStrategy.setEdges(loopNode.getEdges());

        //要添加的变量
        String iterableVarCode = loopNode.getIterableVariable();

        Map<String, Object> variablesToAdd = new HashMap<>();
        variablesToAdd.put(iterableVarCode, iterableValue);
        variablesToAdd.put(LoopNode.ITEM_INDEX_KEY, index);

        //构建结果
        NodeExecuteResult result = new NodeExecuteResult();
        result.setFlowStrategy(flowStrategy);
        result.setVariablesToAdd(variablesToAdd);

        return result;
    }


    private NodeExecuteResult continueNextNode(LoopNode loopNode) {
        //默认策略
        DefaultFlowStrategy flowStrategy = new DefaultFlowStrategy();
        flowStrategy.setNodeId(loopNode.getId());
        flowStrategy.setPortId(loopNode.getOutputPorts().get(0));

        //要移除的变量
        List<String> variablesToRemove = new ArrayList<>();
        variablesToRemove.add(loopNode.getIterableVariable());
        variablesToRemove.add(LoopNode.ITEM_INDEX_KEY);

        //构建结果
        NodeExecuteResult result = new NodeExecuteResult();
        result.setFlowStrategy(flowStrategy);
        result.setVariablesToRemove(variablesToRemove);

        return result;
    }

}